//! 查询引擎插件

use std::sync::Arc;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use datafusion::prelude::*;
use datafusion::execution::context::SessionContext;
use arrow::record_batch::RecordBatch;
use data_gateway_kernel::{
    Plugin, PluginMetadata, PluginStatus, Query, QueryResult, Result
};

/// 查询引擎插件配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryEngineConfig {
    /// 是否启用优化器
    pub enable_optimizer: bool,
    /// 最大内存使用（字节）
    pub max_memory_bytes: usize,
    /// 查询超时时间（秒）
    pub query_timeout_seconds: u64,
    /// 是否启用并行执行
    pub enable_parallel_execution: bool,
    /// 工作线程数
    pub worker_threads: usize,
}

/// 查询引擎插件
pub struct QueryEnginePlugin {
    /// 插件元数据
    metadata: PluginMetadata,
    /// 插件状态
    status: PluginStatus,
    /// DataFusion会话上下文
    session_ctx: Option<Arc<SessionContext>>,
    /// 插件配置
    config: QueryEngineConfig,
    /// 查询统计
    query_stats: QueryEngineStats,
}

/// 查询引擎统计
#[derive(Debug, Clone, Default)]
struct QueryEngineStats {
    total_queries: u64,
    successful_queries: u64,
    failed_queries: u64,
    avg_query_time_ms: f64,
    total_rows_processed: u64,
}

impl Default for QueryEngineConfig {
    fn default() -> Self {
        Self {
            enable_optimizer: true,
            max_memory_bytes: 1024 * 1024 * 1024, // 1GB
            query_timeout_seconds: 300, // 5分钟
            enable_parallel_execution: true,
            worker_threads: num_cpus::get(),
        }
    }
}

impl QueryEnginePlugin {
    /// 创建新的查询引擎插件
    pub fn new(config: QueryEngineConfig) -> Self {
        let metadata = PluginMetadata {
            id: "query_engine".to_string(),
            name: "DataFusion查询引擎插件".to_string(),
            version: "1.0.0".to_string(),
            description: "基于Apache DataFusion的高性能SQL查询引擎".to_string(),
            author: "Data Gateway Team".to_string(),
            capabilities: vec!["sql".to_string()],
            config_schema: Some(serde_json::to_value(&config).unwrap_or_default()),
        };
        
        Self {
            metadata,
            status: PluginStatus::Stopped,
            session_ctx: None,
            config,
            query_stats: QueryEngineStats::default(),
        }
    }
    
    /// 创建带默认配置的查询引擎插件
    pub fn with_default_config() -> Self {
        Self::new(QueryEngineConfig::default())
    }
    
    /// 更新配置
    pub fn update_config(&mut self, config: QueryEngineConfig) -> Result<()> {
        if matches!(self.status, PluginStatus::Running) {
            return Err(anyhow::anyhow!("无法在运行时更新配置").into());
        }
        
        self.config = config;
        self.metadata.config_schema = serde_json::to_value(&self.config).unwrap_or_default();
        
        Ok(())
    }
    
    /// 获取配置
    pub fn get_config(&self) -> &QueryEngineConfig {
        &self.config
    }
    
    /// 执行SQL查询
    async fn execute_sql_query(&mut self, sql: &str) -> Result<RecordBatch> {
        let ctx = self.session_ctx.as_ref()
            .ok_or_else(|| anyhow::anyhow!("查询引擎未初始化"))?;
        
        let start_time = std::time::Instant::now();
        
        tracing::debug!("执行SQL查询: {}", sql);
        
        // 解析并执行SQL
        let df = ctx.sql(sql).await
            .map_err(|e| anyhow::anyhow!("SQL解析失败: {}", e))?;
        
        // 收集结果
        let batches = df.collect().await
            .map_err(|e| anyhow::anyhow!("SQL执行失败: {}", e))?;
        
        let duration = start_time.elapsed().as_millis() as f64;
        
        // 合并结果批次
        let result_batch = if batches.is_empty() {
            // 创建空的RecordBatch
            RecordBatch::new_empty(Arc::new(arrow::datatypes::Schema::empty()))
        } else if batches.len() == 1 {
            batches.into_iter().next().unwrap()
        } else {
            // 合并多个批次
            arrow::compute::concat_batches(&batches[0].schema(), &batches)
                .map_err(|e| anyhow::anyhow!("合并查询结果失败: {}", e))?
        };
        
        // 更新统计信息
        self.update_query_stats(duration, true, result_batch.num_rows());
        
        tracing::debug!("SQL查询完成: {} 行, {}ms", result_batch.num_rows(), duration);
        
        Ok(result_batch)
    }
    
    /// 注册表
    pub async fn register_table(&self, name: &str, table: Arc<dyn datafusion::datasource::TableProvider>) -> Result<()> {
        let ctx = self.session_ctx.as_ref()
            .ok_or_else(|| anyhow::anyhow!("查询引擎未初始化"))?;
        
        ctx.register_table(name, table)
            .map_err(|e| anyhow::anyhow!("注册表失败: {}", e))?;
        
        tracing::info!("已注册表: {}", name);
        Ok(())
    }
    
    /// 注册CSV文件
    pub async fn register_csv(&self, name: &str, file_path: &str) -> Result<()> {
        let ctx = self.session_ctx.as_ref()
            .ok_or_else(|| anyhow::anyhow!("查询引擎未初始化"))?;
        
        ctx.register_csv(name, file_path, CsvReadOptions::new()).await
            .map_err(|e| anyhow::anyhow!("注册CSV文件失败: {}", e))?;
        
        tracing::info!("已注册CSV文件: {} -> {}", name, file_path);
        Ok(())
    }
    
    /// 注册Parquet文件
    pub async fn register_parquet(&self, name: &str, file_path: &str) -> Result<()> {
        let ctx = self.session_ctx.as_ref()
            .ok_or_else(|| anyhow::anyhow!("查询引擎未初始化"))?;
        
        ctx.register_parquet(name, file_path, ParquetReadOptions::default()).await
            .map_err(|e| anyhow::anyhow!("注册Parquet文件失败: {}", e))?;
        
        tracing::info!("已注册Parquet文件: {} -> {}", name, file_path);
        Ok(())
    }
    
    /// 获取已注册的表列表
    pub fn list_tables(&self) -> Result<Vec<String>> {
        let ctx = self.session_ctx.as_ref()
            .ok_or_else(|| anyhow::anyhow!("查询引擎未初始化"))?;
        
        let catalog = ctx.catalog("datafusion").unwrap();
        let schema = catalog.schema("public").unwrap();
        
        Ok(schema.table_names())
    }
    
    /// 更新查询统计
    fn update_query_stats(&mut self, duration_ms: f64, success: bool, rows_processed: usize) {
        self.query_stats.total_queries += 1;
        if success {
            self.query_stats.successful_queries += 1;
            self.query_stats.total_rows_processed += rows_processed as u64;
        } else {
            self.query_stats.failed_queries += 1;
        }
        
        // 计算平均查询时间
        let total_time = self.query_stats.avg_query_time_ms * (self.query_stats.total_queries - 1) as f64 + duration_ms;
        self.query_stats.avg_query_time_ms = total_time / self.query_stats.total_queries as f64;
    }
}

#[async_trait]
impl Plugin for QueryEnginePlugin {
    fn metadata(&self) -> &PluginMetadata {
        &self.metadata
    }
    
    fn status(&self) -> PluginStatus {
        self.status.clone()
    }
    
    async fn start(&mut self) -> Result<()> {
        if matches!(self.status, PluginStatus::Running) {
            return Ok(());
        }
        
        tracing::info!("启动查询引擎插件");
        
        // 创建DataFusion会话上下文
        let mut config = SessionConfig::new()
            .with_target_partitions(self.config.worker_threads);
        
        if !self.config.enable_optimizer {
            // 禁用优化器（如果需要）
            // config = config.with_optimizer_rules(vec![]);
        }
        
        let ctx = SessionContext::with_config(config);
        
        self.session_ctx = Some(Arc::new(ctx));
        self.status = PluginStatus::Running;
        
        tracing::info!("查询引擎插件启动成功");
        Ok(())
    }
    
    async fn stop(&mut self) -> Result<()> {
        if matches!(self.status, PluginStatus::Stopped) {
            return Ok(());
        }
        
        tracing::info!("停止查询引擎插件");
        
        self.session_ctx = None;
        self.status = PluginStatus::Stopped;
        
        tracing::info!("查询引擎插件已停止");
        Ok(())
    }
    
    async fn handle_query(&mut self, query: Query) -> Result<QueryResult> {
        if !matches!(self.status, PluginStatus::Running) {
            return Err(anyhow::anyhow!("插件未运行").into());
        }
        
        if query.query_type != "sql" {
            return Err(anyhow::anyhow!("不支持的查询类型").into());
        }
        
        tracing::debug!("查询引擎插件处理查询: {}", query.content);
        
        let record_batch = self.execute_sql_query(&query.content).await?;
        
        // 转换为内核查询结果
        let mut stats = data_gateway_kernel::QueryStats::new();
        stats.set_rows_returned(record_batch.num_rows());
        stats.finish();
        
        let result = QueryResult::success(
            query.id,
            data_gateway_kernel::QueryResultData::Arrow(vec![record_batch]),
            stats,
        );
        
        tracing::debug!("查询引擎查询完成: {} 行", result.as_arrow().map(|a| a.len()).unwrap_or(0));
        
        Ok(result)
    }
    
    fn can_handle_query(&self, query: &Query) -> bool {
        query.query_type == "sql" && matches!(self.status, PluginStatus::Running)
    }
    
    async fn get_health(&self) -> Result<serde_json::Value> {
        let mut health = serde_json::Map::new();
        
        health.insert("plugin_id".to_string(), serde_json::Value::String(self.metadata.id.clone()));
        health.insert("status".to_string(), serde_json::Value::String(format!("{:?}", self.status)));
        health.insert("enable_optimizer".to_string(), serde_json::Value::Bool(self.config.enable_optimizer));
        health.insert("worker_threads".to_string(), serde_json::Value::Number(self.config.worker_threads.into()));
        
        health.insert("total_queries".to_string(), serde_json::Value::Number(self.query_stats.total_queries.into()));
        health.insert("successful_queries".to_string(), serde_json::Value::Number(self.query_stats.successful_queries.into()));
        health.insert("failed_queries".to_string(), serde_json::Value::Number(self.query_stats.failed_queries.into()));
        health.insert("total_rows_processed".to_string(), serde_json::Value::Number(self.query_stats.total_rows_processed.into()));
        health.insert("avg_query_time_ms".to_string(), serde_json::Value::Number(
            serde_json::Number::from_f64(self.query_stats.avg_query_time_ms).unwrap_or_default()
        ));
        
        // 添加已注册的表信息
        if let Ok(tables) = self.list_tables() {
            health.insert("registered_tables".to_string(), serde_json::Value::Array(
                tables.into_iter().map(serde_json::Value::String).collect()
            ));
        }
        
        Ok(serde_json::Value::Object(health))
    }
    
    async fn update_config(&mut self, config: serde_json::Value) -> Result<()> {
        let new_config: QueryEngineConfig = serde_json::from_value(config)
            .map_err(|e| anyhow::anyhow!("配置格式错误: {}", e))?;
        
        self.update_config(new_config)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_query_engine_plugin_creation() {
        let config = QueryEngineConfig::default();
        let plugin = QueryEnginePlugin::new(config);
        
        assert_eq!(plugin.metadata().id, "query_engine");
        assert_eq!(plugin.metadata().name, "DataFusion查询引擎插件");
        assert!(matches!(plugin.status(), PluginStatus::Stopped));
        assert!(plugin.metadata().capabilities.contains(&"sql".to_string()));
    }
    
    #[test]
    fn test_query_engine_config_update() {
        let mut plugin = QueryEnginePlugin::with_default_config();
        
        let new_config = QueryEngineConfig {
            enable_optimizer: false,
            worker_threads: 4,
            ..Default::default()
        };
        
        let result = plugin.update_config(new_config.clone());
        assert!(result.is_ok());
        assert!(!plugin.get_config().enable_optimizer);
        assert_eq!(plugin.get_config().worker_threads, 4);
    }
    
    #[tokio::test]
    async fn test_plugin_lifecycle() {
        let mut plugin = QueryEnginePlugin::with_default_config();
        
        // 初始状态应该是停止的
        assert!(matches!(plugin.status(), PluginStatus::Stopped));
        
        // 启动插件
        let result = plugin.start().await;
        assert!(result.is_ok());
        assert!(matches!(plugin.status(), PluginStatus::Running));
        
        // 停止插件
        let stop_result = plugin.stop().await;
        assert!(stop_result.is_ok());
        assert!(matches!(plugin.status(), PluginStatus::Stopped));
    }
    
    #[tokio::test]
    async fn test_simple_sql_query() {
        let mut plugin = QueryEnginePlugin::with_default_config();
        plugin.start().await.unwrap();
        
        // 测试简单的SQL查询
        let query = Query::sql("SELECT 1 as id, 'test' as name".to_string());
        let result = plugin.handle_query(query).await;
        
        assert!(result.is_ok());
        let query_result = result.unwrap();
        assert!(query_result.is_success());
        
        if let Some(batches) = query_result.as_arrow() {
            assert_eq!(batches.len(), 1);
            assert_eq!(batches[0].num_rows(), 1);
            assert_eq!(batches[0].num_columns(), 2);
        }
    }
}
